import pandas as pd
from autogluon.core import TabularDataset

from MultilabelPredictor import MultilabelPredictor


class AutogluonModel:

    def __init__(self, path):
        self.predictor = None
        self.path = path
        self.df = None
        # 定义特征和目标列
        self.features = None
        self.targets = None

    def save_prepared_data(self, save_path):
        self.df.to_excel(save_path, index=False)

    def fit(self, train_df=None, features=None, targets=None):
        if train_df is None:
            train_df = self.df
        if features is None:
            features = self.features
        if targets is None:
            targets = self.targets
        target_num = len(targets)
        # 创建AutoGluon数据集
        train_data = TabularDataset(train_df[features + targets])
        # 初始化并训练预测模型
        predictor = MultilabelPredictor(
            labels=targets,
            path=self.path,
            problem_types=['regression'] * target_num,
            eval_metrics=['r2'] * target_num
        )
        self.predictor = predictor
        predictor.fit(
            train_data=train_data,
            time_limit=600,  # 10分钟训练时间
            presets='best_quality',  # 使用最佳质量预设
            verbosity=2  # 显示训练日志
        )
        self.predictor = predictor
        return predictor.leaderboards()

    def predict(self, df, features=None, save_path=None):
        if features is None:
            features = self.features
        predictor = MultilabelPredictor.load(self.path)
        self.predictor = predictor
        result_df = predictor.predict(df[features])
        result_df = pd.concat([df, result_df], axis=1)
        if save_path is not None:
            result_df.to_excel(save_path, index=False)
        return result_df
